import json
import re
from collections import defaultdict
from operator import itemgetter, attrgetter
from typing import Union, Type, Callable, List, Optional, Any

from autogen import GroupChat, GroupChatManager, ConversableAgent, AssistantAgent, UserProxyAgent, Agent
from utils.configuration import load_config, set_environment, default, exists
import string
from string import Template
from typing import Annotated, Literal

Operator = Literal["+", "-", "*", "/"]
config = load_config('./env.yaml')
llm_config = {"config_list": [{"model": config['llm']['model'],
                               "api_key": 'sk-ypDOrct5sWOZEVCUYWmNqK8rsGXcJ6svajXmNrPdEeU5tshv',
                               'base_url': config['llm']['base_url'],
                               }]}

# llm_config = {"config_list": [{"model": 'gpt-3.5-turbo',
#                                "api_key": 'sk-ypDOrct5sWOZEVCUYWmNqK8rsGXcJ6svajXmNrPdEeU5tshv',
#                                'base_url': 'https://api.chatanywhere.tech' #config['llm']['base_url']
#                                }]}

agent = AssistantAgent(
    name='assi',
    llm_config=llm_config,
    system_message="""
    You are a helpful AI assistant.
    A calculator useful for calculation which you can call. 
    You must call calculation tool to finish calculation.
    Return 'TERMINATE' when the task is done. """
)


@agent.register_for_llm(description="A calculator useful for calculation", name="calculator")
def calculator(a: int, b: int, operator: Annotated[Operator, "operator"]) -> int:
    if operator == "+":
        return a + b
    elif operator == "-":
        return a - b
    elif operator == "*":
        return a * b
    elif operator == "/":
        return int(a / b)
    else:
        raise ValueError("Invalid operator")


user_proxy_agent = ConversableAgent(
    name="User",
    llm_config=False,
    is_termination_msg=lambda msg: msg.get("content") is not None and "TERMINATE" in msg["content"],
    human_input_mode="NEVER",
)


def appending_tools_prompt(agent: ConversableAgent):
    system_prompt = agent.system_message
    function_map = agent.function_map
    tools = ",".join(function_map.keys())
    desc = ""
    for name,func in function_map.items():
        desc += f"{name}: {func._description}\n"

    prompts = Template("""
            # 工具
        
        ## 你拥有如下工具：
        
        $tool_descs
        
        ## 你可以在回复中插入零次、一次或多次以下命令以调用工具：
        
        如果需要使用function请回复 "function_call"
        FUNCTION: 工具名称，必须是[ $tool_name ]之一。
        ARGS: 工具输入
        RESULT: 工具结果，需将图片用![](url)渲染出来。
        RETURN: 根据工具结果进行回复 
    """)
    prompt = prompts.substitute(tool_descs=desc, tool_name=tools)
    agent.update_system_message(prompt)



agent.register_function({"calculator": calculator})
appending_tools_prompt(agent)
# user_proxy_agent.register_for_execution(name="calculator")(calculator)
#
chat_result = user_proxy_agent.initiate_chat(agent, message="What is (44232 + 13312 / (232 - 32)) * 5?")
print(chat_result)
